package path_finding;

import java.util.*;

public class AStar implements PathFinding {
    public boolean allowDiagonal = true;
    public Heuristic heuristic = (fromX, fromY, toX, toY) -> {
        int dx = Math.abs(toX - fromX);
        int dy = Math.abs(toY - fromY);
        return Math.min(dx, dy) * 14 + Math.abs(dx - dy) * 10;
    };
    public XyTrans xyTrans;

    @Override
    public List<Integer> findPath(int[][] grid, int startX, int startY, int endX, int endY, int visited[][]) {
        Queue<Node> openSet = new PriorityQueue<>(Comparator.comparingInt(o -> o.f));
        Map<Integer, Node> allNodes = new HashMap<>();
        Map<Integer, Node> closeNodes = new HashMap<>();
        Node startNode = new Node(xyTrans.toId(startX, startY), 0, heuristic.calculate(startX, startY, endX, endY));
        openSet.add(startNode);
        allNodes.put(startNode.id, startNode);
        int endNodeId = xyTrans.toId(endX, endY);
        Node endNode = null;

        while (!openSet.isEmpty()) {
            Node node = openSet.poll();
            allNodes.remove(node.id);
            closeNodes.put(node.id, node);

            if (node.id == endNodeId) {
                endNode = node;
                break;
            }

            int []xy = xyTrans.toXy(node.id);
            List<Integer> neighbors = new ArrayList<>(8);
            getNeighbors(node.id, grid, neighbors);

            for (int neighbor : neighbors) {
                if (closeNodes.containsKey(neighbor)) continue;
                int []neighborXY = xyTrans.toXy(neighbor);
                visited[neighborXY[0]][neighborXY[1]] = 1;
                int cost = 10;
                if (Math.abs(neighborXY[0] - xy[0]) == 1 && Math.abs(neighborXY[1] - xy[1]) == 1) {
                    cost = 14;
                }
                Node neighborNode = new Node(neighbor, node.g + cost, heuristic.calculate(neighborXY[0], neighborXY[1],
                        endX, endY), node);
                if (!allNodes.containsKey(neighbor)) {
                    openSet.add(neighborNode);
                    allNodes.put(neighbor, neighborNode);
                } else {
                    Node existNode = allNodes.get(neighbor);
                    if (neighborNode.f < existNode.f) {
                        openSet.remove(existNode);
                        openSet.add(neighborNode);
                        allNodes.put(neighbor, neighborNode);
                    }
                }
            }
        }

        List findPath = new LinkedList();
        while (endNode != null) {
            findPath.add(endNode.id);
            endNode = endNode.parent;
        }
        Collections.reverse(findPath);
        return findPath;
    }

    private void getNeighbors(int nodeId, int [][] grid, List<Integer> neighbors) {
        int[] dx = {-1, 1, 0, 0};
        int[] dy = {0, 0, -1, 1};

        // 对角线方向
        if (allowDiagonal) {
            int[] diagonalDx = {-1, -1, 1, 1};
            int[] diagonalDy = {-1, 1, -1, 1};
            for (int i = 0; i < 4; i++) {
                addNeighborIfValid(nodeId, diagonalDx[i], diagonalDy[i], grid, neighbors);
            }
        }

        // 四方向
        for (int i = 0; i < 4; i++) {
            addNeighborIfValid(nodeId, dx[i], dy[i], grid, neighbors);
        }
    }

    private void addNeighborIfValid(int nodeId, int dx, int dy, int[][] grid, List<Integer> neighbors) {
        int []xy = xyTrans.toXy(nodeId);
        int newX = xy[0] + dx;
        int newY = xy[1] + dy;

        if (newX >= 0 && newX < grid.length &&
                newY >= 0 && newY < grid[0].length &&
                grid[newX][newY] == 0) { // 0 表示可通行
            neighbors.add(xyTrans.toId(newX, newY));
        }
    }

    @Override
    public void setAllowDiagonal(boolean allow) {
        this.allowDiagonal = allow;
    }

    @Override
    public void setHeuristic(Heuristic heuristic) {
        this.heuristic = heuristic;
    }

    @Override
    public void setXyTrans(XyTrans xyTrans) {
        this.xyTrans = xyTrans;
    }

    private static class Node {
        public int id;
        public Node parent;
        public int f;
        public int g;
        public int h;

        public Node(int id, int g, int h) {
            this.id = id;
            this.f = g + h;
            this.g = g;
            this.h = h;
        }

        public Node(int id, int g, int h, Node parent) {
            this.id = id;
            this.f = g + h;
            this.parent = parent;
            this.g = g;
            this.h = h;
        }
    }
}
